import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.axes import Axes
from matplotlib.figure import Figure
from scipy.stats import poisson
import matplotlib
import numpy.random as npr
matplotlib.use("TkAgg")

# 固定 lambda
lam = 5

# 不同的 n 和 p = lambda/n
n_values = [20, 50, 100, 200, 500, 1000, 2000,3000,4000,5000]
p_values = [lam / n for n in n_values]

# X 取值范围
x = np.arange(0, 20)

fig:Figure
ax:Axes
# 创建图像
fig, ax = plt.subplots(figsize=(8, 5))

# 初始化柱状图和泊松曲线
bars = ax.bar(x=x, height=np.zeros_like(x), alpha=0.6, label="Binomial Samples")
poisson_line, = ax.plot(x, poisson.pmf(x, lam), marker='o', linestyle='-', linewidth=2, label=f"Poisson PMF (λ={lam})")

# 设置图例和标题
ax.set_ylim(bottom=0, top=0.2)
ax.set_title(label="二项分布逼近泊松分布（动画）")
ax.set_xlabel("k")
ax.set_ylabel("Probability")
ax.legend()


# 更新函数
def update(frame):
    n = n_values[frame]
    p = p_values[frame]

    # 生成样本
    samples = npr.binomial(n=n, p=p, size=50000)

    # 统计频率
    counts, _ = np.histogram(a=samples, bins=np.arange(21), density=True)

    # 更新柱状图
    for bar, h in zip(bars, counts):
        bar.set_height(h)

    ax.set_title(f"n={n}, p={p:.4f}    (np = λ ≈ {lam})")
    return bars, poisson_line


# 创建动画
ani = FuncAnimation(fig=fig, func=update, frames=len(n_values), interval=2000, repeat=True)

plt.show()
